import re
import json



def remove_comments(code: str) -> str:
    """
    Removes comments from Java, JavaScript, and Python code.
    
    Args:
        code (str): The input code string.
        
    Returns:
        str: The code string with comments removed.
    """
    # Regex for single-line comments and multi-line comments
    java_js_single_line = r'//.*?$'  # Single-line comments for Java/JavaScript
    python_single_line = r'#.*?$'   # Single-line comments for Python
    multi_line = r'/\*.*?\*/'       # Multi-line comments for Java/JavaScript
    
    # Combine regex patterns
    combined_pattern = f"({java_js_single_line}|{python_single_line}|{multi_line})"
    
    # Remove comments using regex with multiline and dotall flags
    no_comments = re.sub(combined_pattern, '', code, flags=re.MULTILINE | re.DOTALL)
    return no_comments


def remove_comments_file(filename):
    with open(f'/data/xingshiyou-slurm/codeLLM/data/answer/{filename}', 'r')as f:
        data_list = json.load(f)
        res = [[remove_comments(data[0])] for data in data_list]

    with open(f'/data/xingshiyou-slurm/codeLLM/data/answer/remove_comments/{filename}', 'w')as f:
        json.dump(res, f, ensure_ascii=False, indent=4)


if __name__ == '__main__':
    for filename in [f'B_maxlen_10000_steps_{index}00.json' for index in range(1,18)]:
        remove_comments_file(filename)
